package com.asmtest.visitor;

import com.util.description.ArrayTypeDesc;
import com.util.description.MethodSignature;
import com.util.description.ParameterizedTypeDesc;
import com.util.description.PrimitiveTypeDesc;
import com.util.description.TypeDesc;
import com.util.description.TypeVariableTypeDesc;
import com.util.description.WildcardTypeDesc;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.signature.SignatureVisitor;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.function.Consumer;

/**
 * @author yutianhong
 * @version 1.0
 * @since 2024/01/01 15:01
 */
public class CustomMethodSignatureVisitor extends BaseSignatureVisitor {

    private final MethodSignature methodSignature = new MethodSignature();

    // fields for formal type parameters
    private String className = null;
    private List<TypeDesc> upperBounds = null;

    @Override
    public void visitFormalTypeParameter(String name) {
        if (className == null) {
            upperBounds = new ArrayList<>();
        } else {
            methodSignature.getTypeVariables().add(TypeVariableTypeDesc.of(className, upperBounds));
        }
        className = name;
    }

    @Override
    public SignatureVisitor visitClassBound() {
        return new TypeVisitor(typeDesc -> upperBounds.add(typeDesc));
    }

    @Override
    public SignatureVisitor visitInterfaceBound() {
        return new TypeVisitor(typeDesc -> upperBounds.add(typeDesc));
    }

    @Override
    public SignatureVisitor visitParameterType() {
        if (className != null) {
            methodSignature.getTypeVariables().add(TypeVariableTypeDesc.of(className, upperBounds));
            className = null;
        }
        return new TypeVisitor(typeDesc -> methodSignature.getParamTypes().add(typeDesc));
    }

    @Override
    public SignatureVisitor visitReturnType() {
        return new TypeVisitor(methodSignature::setReturnType);
    }

    @Override
    public SignatureVisitor visitExceptionType() {
        return new TypeVisitor(typeDesc -> methodSignature.getExceptions().add(typeDesc));
    }

    public MethodSignature getMethodSignature() {
        return methodSignature;
    }
}

class BaseSignatureVisitor extends SignatureVisitor {
    public BaseSignatureVisitor() {
        super(Opcodes.ASM9);
    }
}

class TypeVisitor extends BaseSignatureVisitor {
    boolean array;
    boolean typeVariable;
    String className;
    boolean wildcard;
    boolean primitive;
    boolean parameterized;
    char boundDesc;
    List<TypeDesc> types;
    final Consumer<TypeDesc> consumer;

    public TypeVisitor(Consumer<TypeDesc> consumer) {
        this.consumer = consumer;
        init();
    }

    @Override
    public SignatureVisitor visitArrayType() {
        array = true;
        return this;
    }

    @Override
    public void visitBaseType(char descriptor) {
        // include "void" type
        primitive = true;
        className = String.valueOf(descriptor);
        visitEnd();
    }

    @Override
    public void visitClassType(String name) {
        className = name;
    }

    @Override
    public void visitTypeVariable(String name) {
        typeVariable = true;
        className = name;
        visitEnd();
    }

    @Override
    public void visitTypeArgument() {
        parameterized = true;
        types.add(WildcardTypeDesc.of(Collections.emptyList(), Collections.emptyList()));
    }

    @Override
    public SignatureVisitor visitTypeArgument(char wildcard) {
        parameterized = true;
        TypeVisitor typeArgVisitor = new TypeVisitor(typeDesc -> types.add(typeDesc));
        if (wildcard != INSTANCEOF) {
            typeArgVisitor.wildcard = true;
            typeArgVisitor.boundDesc = wildcard;
        }
        return typeArgVisitor;
    }

    @Override
    public void visitEnd() {
        TypeDesc typeDesc;
        if (parameterized) {
            typeDesc = ParameterizedTypeDesc.of(className, types);
        } else if (typeVariable) {
            typeDesc = TypeVariableTypeDesc.of(className, types);
        } else if (primitive) {
            typeDesc = PrimitiveTypeDesc.getByClassName(className);
        } else {
            typeDesc = TypeDesc.of(className);
        }
        if (array) {
            typeDesc = ArrayTypeDesc.of(typeDesc);
        } else if (wildcard) {
            List<TypeDesc> upperBounds = boundDesc == EXTENDS ? Collections.singletonList(typeDesc) : Collections.emptyList();
            List<TypeDesc> lowerBounds = boundDesc == SUPER ? Collections.singletonList(typeDesc) : Collections.emptyList();
            typeDesc = WildcardTypeDesc.of(upperBounds, lowerBounds);
        }
        consumer.accept(typeDesc);
        init();
    }

    private void init() {
        // init
        array = false;
        typeVariable = false;
        className = "";
        wildcard = false;
        primitive = false;
        parameterized = false;
        boundDesc = 0;
        types = new ArrayList<>();
    }
}
